#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time
import random
import torch

from distance import euclidean_squared_distance, euclidean_squared_distance2, cosine_distance, cosine_distance2, cosine_distance_ex

sys.path.append('.')

"""
    手工通过种子人脸将人脸合并
"""


video_base_paths = [
    # r'/rootfs/media/kasim/DataSet/RecordEx',
    # r'/rootfs/media/kasim/DataSet/Record',
    # r'/rootfs/media/kasim/DataSet/Record2',
    r'/rootfs/media/yery/Dataset/RecordEx',
    r'/rootfs/media/yery/Dataset/Record',
    r'/rootfs/media/yery/Dataset/Record2',
]

DISPLAY_WIDTH = 1280
DISPLAY_HEIGHT = 720

h_image_count = 10
v_image_count = 10
img_width = 200
img_height = 200
# img_width = 96
# img_height = 96

batch_size = 32

ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()
current_mouse_xy = None


def get_xy_bbox_idx(xy):
    global image_bboxes
    bbox_idx = None
    if xy is None:
        return bbox_idx
    x, y = xy
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes, current_mouse_xy

    current_mouse_xy = (x, y)

    if draw_image is not None:
        if event == cv2.EVENT_LBUTTONDOWN:
            # print(x, y)
            bbox_idx = get_xy_bbox_idx(current_mouse_xy)
            if bbox_idx is not None:
                if bbox_idx not in select_image_bbox_idxes:
                    select_image_bbox_idxes.add(bbox_idx)
                else:
                    select_image_bbox_idxes.remove(bbox_idx)


def image_resize(img, img_width, img_height, mode=0):
    if mode == 0:
        img = cv2.resize(img, (img_width, img_height))
    elif mode == 1:
        w, h = img.shape[1], img.shape[0]
        wscale = img_width / w
        hscale = img_height / h
        scale = min(wscale, hscale)
        w = int(round(w * scale))
        h = int(round(h * scale))
        image = cv2.resize(img, (w, h))
        img = np.full((img_height, img_width, image.shape[2]), fill_value=0, dtype=np.uint8)
        x = (img_width - w) // 2
        y = (img_height - h) // 2
        img[y:y + h, x:x + w, :] = image
    return img


def show_images(image_dir, file_names, dis_text='', win_name='image', wait=False, h_image_count=10, v_image_count=10, display_mode=0):
    display_image_count = h_image_count * v_image_count
    files = []
    total_image_count = len(file_names)
    last_idx = total_image_count - 1
    # if total_image_count > 0:
    #     print('{}: {}'.format(win_name, file_names[0]))
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            if file_name is None:
                img = np.zeros((img_height, img_width, 3), dtype=np.uint8)
            else:
                file_path = os.path.join(image_dir, file_name.strip())
                img = cv2.imread(file_path)
                img = image_resize(img, img_width, img_height, mode=display_mode)
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (2, 12),
                    cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return
        else:
            break


def show_images_ex(image_dir, file_names_list, dis_text='', win_name='image', wait=False, h_image_count=20, v_image_count=10, display_mode=0):
    file_counts = [len(file_names) for file_names in file_names_list]
    max_file_counts = max(file_counts)
    if max_file_counts < h_image_count:
        h_image_count = max_file_counts

    file_names = []
    for file_count, _file_names in zip(file_counts, file_names_list):
        if file_count >= h_image_count:
            file_names.extend(_file_names[:h_image_count])
        else:
            file_names.extend(_file_names)
            for i in range(h_image_count-file_count):
                file_names.append(None)

    show_images(image_dir, file_names, dis_text=dis_text, win_name=win_name, wait=wait, h_image_count=h_image_count, v_image_count=v_image_count, display_mode=display_mode)


def display_video(video_path, frame_num=None, win_name='video'):
    cap = cv2.VideoCapture(video_path)
    is_pause = False
    frame_id = 0
    while True:
        grabbed, image = cap.read()

        if not grabbed:
            break

        cv2.putText(image, 'Frame: {}'.format(frame_id+1), (32, 32), cv2.FONT_HERSHEY_COMPLEX, 1.0, (255, 0, 0))

        if (image.shape[0] != DISPLAY_HEIGHT) or (image.shape[1] != DISPLAY_WIDTH):
            image = cv2.resize(image, (DISPLAY_WIDTH, DISPLAY_HEIGHT))
        cv2.imshow(win_name, image)
        if ((frame_num is not None) and (frame_id == frame_num)) or is_pause:
            key = cv2.waitKey()
        else:
            key = cv2.waitKey(40)
        if key == ord('g') or key == ord('G'):
            is_pause = False
        elif key == ord('z') or key == ord('Z'):
            is_pause = True
        elif key == ord('o') or key == ord('O'):
            break

        frame_id += 1

    cap.release()
    cv2.destroyWindow(win_name)
    print(video_path)


def play_video(video_name, frame_num=None, win_name='video'):
    for video_base_path in video_base_paths:
        video_path = os.path.join(video_base_path, video_name)
        if not os.path.exists(video_path):
            continue
        display_video(video_path, frame_num=frame_num, win_name=win_name)


def save_process(cur_xiaoqu_serial_num, xiaoqu_faceid_file_feature_faceids, faceid_file_list_path, cur_id_idx_file_path):
    with open(faceid_file_list_path, 'w', encoding='utf-8') as file:
        for xiaoqu, faceid_file_feature_faceids in xiaoqu_faceid_file_feature_faceids.items():
            for faceid, file_feature_faceids in faceid_file_feature_faceids.items():
                if file_feature_faceids is None:
                    continue
                faceids = file_feature_faceids[2]
                serial_num = file_feature_faceids[3]
                same_faceid_nums_info = xiaoqu + ':' + (','.join(faceids)) + ':' + str(serial_num) + '\n'
                file.write(same_faceid_nums_info)
    os.system('chmod a+wr {}'.format(faceid_file_list_path))

    with open(cur_id_idx_file_path, 'w', encoding='utf-8') as file:
        for xiaoqu, next_serial_num in cur_xiaoqu_serial_num.items():
            process_info = '{}:{}\n'.format(xiaoqu, next_serial_num)
            file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_id_idx_file_path))


def read_features(file_list, feature_file, feat_dim, mode=0):
    xiaoqu_faceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            if mode == 0:
                # 150100414a54443452064d2742ce3600 fbea3e79-e936-4307-9f6a-c4dec71e2f92 gd/zs/nqjd/dxjyjy/150100414a54443452064d2742ce3600/video/20190710_133107/000040_000.jpg 1
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                if len(info) > 1:
                    xiaoqu = '/'.join(info[:-4])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 1:
                # gd-zs-sqqjd-dxssjyeq-150100414a54443452064d2742c73600-video-20190913_225440/0001/0001_F0001B00.jpg
                line = line.strip()
                lines = line.split('/')
                file_name = line
                info = lines[0].split('-')
                if len(info) > 1:
                    xiaoqu = '-'.join(info[:-3])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 2:
                # c187ca1c31293522b0331c0ca37d4eee/dcb4b86dd6c6e155de90d3311674d754.jpg 21,45,164,236,1,1024
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                xiaoqu = 'NONE'
                faceid = info[0]

            # feature = np.fromfile(feat_file, dtype=np.float32, count=feat_dim).reshape(1, feat_dim)
            feature = np.fromfile(feat_file, dtype=np.float32, count=feat_dim)

            faceid_file_features = xiaoqu_faceid_file_features.get(xiaoqu, None)
            if faceid_file_features is None:
               files = [file_name]
               features = [feature]
               file_features = [files, features]
               xiaoqu_faceid_file_features[xiaoqu] = {faceid: file_features}
            else:
                file_features = faceid_file_features.get(faceid, None)
                if file_features is None:
                    files = [file_name]
                    features = [feature]
                    file_features = [files, features]
                    faceid_file_features[faceid] = file_features
                else:
                    file_features[0].append(file_name)
                    file_features[1].append(feature)
        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()

    return xiaoqu_faceid_file_features


def read_file_list(file_list, mode=0):
    xiaoqu_faceid_files = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            if mode == 0:
                # 150100414a54443452064d2742ce3600 fbea3e79-e936-4307-9f6a-c4dec71e2f92 gd/zs/nqjd/dxjyjy/150100414a54443452064d2742ce3600/video/20190710_133107/000040_000.jpg 1
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                if len(info) > 1:
                    xiaoqu = '/'.join(info[:-4])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 1:
                # gd-zs-sqqjd-dxssjyeq-150100414a54443452064d2742c73600-video-20190913_225440/0001/0001_F0001B00.jpg
                line = line.strip()
                lines = line.split('/')
                file_name = line
                info = lines[0].split('-')
                if len(info) > 1:
                    xiaoqu = '-'.join(info[:-3])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 2:
                # c187ca1c31293522b0331c0ca37d4eee/dcb4b86dd6c6e155de90d3311674d754.jpg 21,45,164,236,1,1024
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                xiaoqu = 'NONE'
                faceid = info[0]

            faceid_files = xiaoqu_faceid_files.get(xiaoqu, None)
            if faceid_files is None:
               files = [file_name]
               xiaoqu_faceid_files[xiaoqu] = {faceid: files}
            else:
                files = faceid_files.get(faceid, None)
                if files is None:
                    files = [file_name]
                    faceid_files[faceid] = files
                else:
                    files.append(file_name)
        print('Read File Count: {}'.format(file_count))

    return xiaoqu_faceid_files


def read_faceid_list(file_list, image_list):
    xiaoqu_faceids = {}
    file_count = 0
    image_file = open(image_list, 'r')
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read FaceID Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            faceid_sub = line.strip()

            file_name = image_file.readline().strip()
            info = file_name.split('/')
            if len(info) > 1:
                xiaoqu = '/'.join(info[:-4])
            else:
                xiaoqu = 'NONE'

            faceids = xiaoqu_faceids.get(xiaoqu, None)
            if faceids is None:
                xiaoqu_faceids[xiaoqu] = [faceid_sub]
            else:
                faceids.append(faceid_sub)
        print('Read FaceID Count: {}'.format(file_count))
    image_file.close()
    return xiaoqu_faceids


def read_faceid_list2(file_list, mode=0):
    xiaoqu_faceids = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read FaceID Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            if mode == 0:
                # 150100414a54443452064d2742ce3600 fbea3e79-e936-4307-9f6a-c4dec71e2f92 gd/zs/nqjd/dxjyjy/150100414a54443452064d2742ce3600/video/20190710_133107/000040_000.jpg 1
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                if len(info) > 1:
                    xiaoqu = '/'.join(info[:-4])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 1:
                # gd-zs-sqqjd-dxssjyeq-150100414a54443452064d2742c73600-video-20190913_225440/0001/0001_F0001B00.jpg
                line = line.strip()
                lines = line.split('/')
                file_name = line
                info = lines[0].split('-')
                if len(info) > 1:
                    xiaoqu = '-'.join(info[:-3])
                else:
                    xiaoqu = 'NONE'
                faceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 2:
                # c187ca1c31293522b0331c0ca37d4eee/dcb4b86dd6c6e155de90d3311674d754.jpg 21,45,164,236,1,1024
                lines = line.strip().split()
                file_name = lines[-2].strip()
                info = file_name.split('/')
                xiaoqu = 'NONE'
                faceid = info[0]

            faceids = xiaoqu_faceids.get(xiaoqu, None)
            if faceids is None:
                xiaoqu_faceids[xiaoqu] = {faceid}
            else:
                faceids.add(faceid)
        print('Read FaceID Count: {}'.format(file_count))

    return xiaoqu_faceids


def read_merge_faceid_list(file_list):
    xiaoqu_faceids = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read FaceID Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # NONE:c651b1e5-fdbc-4447-9a6a-9213d7729faa_ac638be4-0e9a-42e8-a29b-c730a41aac50,752524c4-e540-45e2-8388-55aab4d72150_94d61342-0841-432a-a8f2-5bb6cf5c9a4f:0
            lines = line.strip().split(':')
            xiaoqu = lines[0]
            _faceids = lines[1].split(',')
            serial_num = int(lines[2])

            faceids_map = xiaoqu_faceids.get(xiaoqu, None)
            if faceids_map is None:
                xiaoqu_faceids[xiaoqu] = OrderedDict({_faceids[0]: [_faceids, serial_num]})
            else:
                faceids_map[_faceids[0]] = [_faceids, serial_num]
        print('Read FaceID Count: {}'.format(file_count))
    return xiaoqu_faceids


def get_topk(faceid, faceid_file_feature_faceids, topk):
    cur_files = []
    faceids = []
    files = []
    mindists = []
    file_feature_faceids = faceid_file_feature_faceids[faceid]
    _cur_files = file_feature_faceids[0]
    features = file_feature_faceids[1]
    for _faceid, file_feature_faceids in faceid_file_feature_faceids.items():
        if (faceid == _faceid) or (file_feature_faceids is None):
            continue
        serial_num = file_feature_faceids[3]
        if serial_num < 0:
            continue
        _files = file_feature_faceids[0]
        _features = file_feature_faceids[1]
        distmat = cosine_distance2(features, _features)
        mindist_idx = distmat.argmin()
        mindist = (distmat.view(-1)[mindist_idx]).cpu()
        mindist_idx = int(mindist_idx.cpu())
        row_mindist_idx = mindist_idx // _features.size(0)
        mindist_idx = mindist_idx % _features.size(0)
        faceids.append(_faceid)
        files.append(_files[mindist_idx])
        mindists.append(mindist)
        cur_files.append(_cur_files[row_mindist_idx])
    if len(faceids) < 1:
        return [], [], [], []
    mindists = torch.stack(mindists)
    mindist_idx = mindists.argsort(descending=False)
    if topk > mindist_idx.size(0):
        topk = mindist_idx.size(0)
    mindists = mindists[mindist_idx]
    topk_score = 1.0 - mindists[0:topk]
    topk_idx = mindist_idx[0:topk]

    topk_faceids = []
    topk_files = []
    topk_cur_files = []
    for idx in topk_idx:
        topk_faceids.append(faceids[idx])
        topk_files.append(files[idx])
        topk_cur_files.append(cur_files[idx])
    return topk_cur_files, topk_faceids, topk_files, topk_score.tolist()


def merge_faceids_func(faceid, topk_faceids, select_idxes, faceid_file_feature_faceids, faceid_pairs=None):
    if faceid_pairs is not None:
        topk_count = len(topk_faceids)
        for idx in range(topk_count):
            if idx not in select_idxes:
                no_select_faceid = topk_faceids[idx]
                if faceid < no_select_faceid:
                    faceid_pairs.add(faceid + no_select_faceid)
                else:
                    faceid_pairs.add(no_select_faceid + faceid)

    if len(select_idxes) < 1:
        return False
    file_feature_faceids = faceid_file_feature_faceids.get(faceid, None)
    if file_feature_faceids is None:
        return False
    files = file_feature_faceids[0]
    features = file_feature_faceids[1]
    faceids = file_feature_faceids[2]
    features_list = [features]
    has_change = False
    for idx in select_idxes:
        select_faceid = topk_faceids[idx]
        select_file_feature_faceids = faceid_file_feature_faceids.get(select_faceid, None)
        if select_file_feature_faceids is None:
            continue
        if select_file_feature_faceids[3] < 0:
            continue
        sfiles = select_file_feature_faceids[0]
        sfeatures = select_file_feature_faceids[1]
        sfaceids = select_file_feature_faceids[2]
        if len(sfiles) > 0:
            files.extend(sfiles)
            faceids.extend(sfaceids)
            features_list.append(sfeatures)
            has_change = True
        faceid_file_feature_faceids[select_faceid] = None

    if has_change:
        file_feature_faceids[1] = torch.cat(features_list, 0)
    else:
        return False
    return True


def read_select_faceid_list(faceid_list):
    faceids_set = set()
    with open(faceid_list, 'r') as file:
        for line in file.readlines():
            faceid = line.strip().replace(' ', '_', 1)
            faceids_set.add(faceid)
    return faceids_set


def read_device_unit(device_list):
    device_xiaoqu_unit_map = {}
    unit_id = 1
    with open(device_list, 'r', encoding='utf-8') as file:
        for line in file.readlines():
            # 星光礼寓;150100414a54443452067fa1d4c16600,150100414a5444345203bcd0428ab500;5(幢/座) 1(单元) 2梯 A面,5(幢/座) 1(单元) 1梯 A面
            lines = line.split(';')
            xiaoqu = lines[0]
            device_ids = lines[1].split(',')
            for device_id in device_ids:
                device_xiaoqu_unit_map[device_id] = [xiaoqu, unit_id]
            unit_id += 1
    return device_xiaoqu_unit_map


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes
    parser = argparse.ArgumentParser(description="Select Face")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage", help="path to image file", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter7/FilterSelectFace.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter7/FilterSelectFeature.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--typical_file_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter7/FaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--typical_feature_file", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter7/FaceRecFeatSel.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--faceid_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/AttrDiffFace.0.4_0.txt", help="path to need process faceid list file", type=str
    )
    parser.add_argument(
        "--faceid_image_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter6/AttrDiffFaceImage.0.4_0.txt", help="path to need process faceid list file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter9", help="path to output", type=str
    )
    parser.add_argument(
        "--select_faceid_list", default="/rootfs/media/kasim/Data1/data/VideoFaceImage/filter8/FaceIdList.txt", help="path to select face id list", type=str
    )
    parser.add_argument(
        "--device_list", help="same unint devoce list", type=str
    )
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=3*60, help='save time, second ')
    parser.add_argument("--feature_dim", default=256, help="feature dim", type=int)
    parser.add_argument('--max_display_count', type=int, default=8, help='max display count')
    parser.add_argument('--topk', type=int, default=5, help='topk similarity face')
    parser.add_argument('--score_thr', type=float, default=0.6, help='face merge min score thr')
    parser.add_argument("--use_select_faceid_list", action='store_true', help="use select faceid list")
    parser.add_argument('--mode', type=int, default=0, help='mode, 0 is read from faceid_list and faceid_image_list, 1 is read from typical_file_list')
    parser.add_argument('--file_list_mode', type=int, default=0, help='file list mode')
    parser.add_argument('--display_mode', type=int, default=0, help='display mode')
    args = parser.parse_args()

    random.seed(123456)

    DEVICE = 'cuda:0'

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    # display_image_count = h_image_count * v_image_count

    if args.device_list is None:
        device_xiaoqu_unit_map = {}
    else:
        device_xiaoqu_unit_map = read_device_unit(args.device_list)

    if args.use_select_faceid_list:
        select_faceids_set = read_select_faceid_list(args.select_faceid_list)
        process_select_faceid_count = 0

    # xiaoqu_faceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim, args.file_list_mode)
    xiaoqu_faceid_file_features = read_features(args.typical_file_list, args.typical_feature_file, args.feature_dim, args.file_list_mode)

    xiaoqu_faceid_files = read_file_list(args.file_list, args.file_list_mode)

    merge_xiaoqu_faceid_file_feature_faceids = OrderedDict()
    faceid_file_list_path = os.path.join(output_dir, 'SelectMergeFaces.txt')
    if not os.path.exists(faceid_file_list_path):
        if args.mode == 0:
            xiaoqu_faceids = read_faceid_list(args.faceid_list, args.faceid_image_list)
        else:
            xiaoqu_faceids = read_faceid_list2(args.typical_file_list, args.file_list_mode)
        xiaoqus = list(xiaoqu_faceids.keys())
        xiaoqus.sort()
        for xiaoqu in xiaoqus:
            faceids = xiaoqu_faceids[xiaoqu]
            faceid_file_features = xiaoqu_faceid_file_features.get(xiaoqu, None)
            if faceid_file_features is None:
                continue
            merge_faceid_file_feature_faceids = OrderedDict()
            faceids = list(faceids)
            faceids.sort()
            serial_num = 0
            for faceid in faceids:
                file_features = faceid_file_features.get(faceid, None)
                if file_features is None:
                    continue
                if len(file_features) > 0:
                    merge_faceid_file_feature_faceids[faceid] = [file_features[0], file_features[1], [faceid], serial_num]
                    serial_num += 1
            if len(merge_faceid_file_feature_faceids) > 0:
                merge_xiaoqu_faceid_file_feature_faceids[xiaoqu] = merge_faceid_file_feature_faceids
    else:
        xiaoqu_faceids = read_merge_faceid_list(faceid_file_list_path)
        xiaoqus = list(xiaoqu_faceids.keys())
        xiaoqus.sort()
        for xiaoqu in xiaoqus:
            faceids_maps = xiaoqu_faceids[xiaoqu]
            faceid_file_features = xiaoqu_faceid_file_features.get(xiaoqu, None)
            if faceid_file_features is None:
                continue
            merge_faceid_file_feature_faceids = OrderedDict()
            for faceid, faceid_nums in faceids_maps.items():
                merge_faceids, merge_files, merge_features = [], [], []
                _faceids, serial_num = faceid_nums
                for _faceid in _faceids:
                    file_features = faceid_file_features.get(_faceid, None)
                    if file_features is None:
                        continue
                    if len(file_features) > 0:
                        merge_faceids.append(_faceid)
                        merge_files.extend(file_features[0])
                        merge_features.extend(file_features[1])
                if len(merge_faceids) > 0:
                    merge_faceid_file_feature_faceids[faceid] = [merge_files, merge_features, merge_faceids, serial_num]
            if len(merge_faceid_file_feature_faceids) > 0:
                merge_xiaoqu_faceid_file_feature_faceids[xiaoqu] = merge_faceid_file_feature_faceids

    xiaoqu_faceid_file_features.clear()
    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    cur_xiaoqu_serial_num = {}
    cur_id_idx_file_path = os.path.join(output_dir, 'MergeCurIdIdx.txt')
    if os.path.exists(cur_id_idx_file_path):
        with open(cur_id_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_xiaoqu, cur_serial_num = line.strip().split(':')
                cur_xiaoqu_serial_num[cur_xiaoqu] = int(cur_serial_num)

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectFaceImage')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_select_file_list_path = os.path.join(tmp_output_dir, 'SelectMergeFaces.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, 'MergeCurIdIdx.txt')

    merge_change = False
    select_change = False
    start_save_time = time.time()

    topk = args.topk
    xiaoqu_count = len(merge_xiaoqu_faceid_file_feature_faceids)
    xiaoqu_num = 0
    for xiaoqu, faceid_file_feature_faceids in merge_xiaoqu_faceid_file_feature_faceids.items():
        xiaoqu_num += 1
        print(xiaoqu, 0)
        cur_serial_num = cur_xiaoqu_serial_num.get(xiaoqu, None)
        if cur_serial_num is None:
            cur_xiaoqu_serial_num[xiaoqu] = 0
            cur_serial_num = 0
        print(xiaoqu, 1)
        faceid_files = xiaoqu_faceid_files.get(xiaoqu, None)

        faceid_pairs = set()

        face_count = 0
        for faceid, file_feature_faceids in faceid_file_feature_faceids.items():
            features = file_feature_faceids[1]
            # features = [np.frombuffer(feature, dtype=np.float32, count=args.feature_dim) for feature in features]
            features = np.stack(features)
            # features = np.concatenate(features, axis=0)
            features = torch.from_numpy(features)
            # features.share_memory_()
            file_feature_faceids[1] = features.cuda(device=DEVICE)
            if face_count < file_feature_faceids[3]:
                face_count = file_feature_faceids[3]
        face_count += 1
        for faceid, file_feature_faceids in faceid_file_feature_faceids.items():
            if file_feature_faceids is None:
                continue
            serial_num = file_feature_faceids[3]
            cur_xiaoqu_serial_num[xiaoqu] = serial_num
            if ((cur_serial_num is not None) and (serial_num < cur_serial_num)) or (serial_num < 0):
                continue

            if args.use_select_faceid_list:
                if process_select_faceid_count >= len(select_faceids_set):
                    break
                if faceid in select_faceids_set:
                    process_select_faceid_count += 1
                else:
                    continue

            # file_list, features, face_ids = file_feature_faceids
            print('\n', xiaoqu, faceid, '\n')

            _topk_cur_files, _topk_faceids, _topk_files, _topk_score = get_topk(faceid, faceid_file_feature_faceids, topk)
            topk_cur_files, topk_faceids, topk_files, topk_score = [], [], [], []
            for i, score in enumerate(_topk_score):
                if score < args.score_thr:
                    continue
                _topk_faceid = _topk_faceids[i]
                if faceid < _topk_faceid:
                    topk_faceid_pair = faceid + _topk_faceid
                else:
                    topk_faceid_pair = _topk_faceid + faceid
                if topk_faceid_pair in faceid_pairs:
                    # print(topk_faceid_pair)
                    continue
                topk_cur_files.append(_topk_cur_files[i])
                topk_faceids.append(_topk_faceid)
                topk_files.append(_topk_files[i])
                topk_score.append(score)
            if len(topk_faceids) < 1:
                continue

            topk_cur_file = topk_cur_files[0]
            _names = topk_cur_file.split('/')
            cur_device_id = _names[-4]
            cur_device_xiaoqu, cur_device_unit = device_xiaoqu_unit_map.get(cur_device_id, [None, None])
            if cur_device_xiaoqu is None:
                cur_device_xiaoqu = _names[-5]
                cur_device_unit = cur_device_id

            topk_imgs = []
            x_off = 0
            y_off = 0
            image_bboxes.clear()
            for j, file_name in enumerate(topk_files):
                _names = file_name.split('/')
                _device_id = _names[-4]
                _device_xiaoqu, _device_unit = device_xiaoqu_unit_map.get(_device_id, [None, None])
                if _device_unit is None:
                    _device_xiaoqu = _names[-5]
                    _device_unit = _device_id
                unit = 'Same Unit' if _device_unit == cur_device_unit else 'Not Same Unit'
                print(j, file_name, _device_xiaoqu, unit)
                file_path = os.path.join(image_dir, file_name)
                img = cv2.imread(file_path)
                img = image_resize(img, img_width, img_height, mode=args.display_mode)
                topk_imgs.append(img)
                image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                x_off += img_width

            print('c', topk_cur_file, cur_device_xiaoqu)

            image_count = len(topk_imgs)
            images = [np.concatenate(topk_imgs, axis=1)]

            cur_imgs = []
            for file_name in topk_cur_files:
                file_path = os.path.join(image_dir, file_name)
                img = cv2.imread(file_path)
                img = image_resize(img, img_width, img_height, mode=args.display_mode)
                cur_imgs.append(img)
            images.append(np.concatenate(cur_imgs, axis=1))

            draw_image = np.vstack(images)
            select_image_bbox_idxes.clear()

            file_names_list = []
            for topk_faceid in topk_faceids:
                file_names_list.append(faceid_file_feature_faceids[topk_faceid][0])
            file_names_list.append(file_feature_faceids[0])
            show_images_ex(image_dir, file_names_list, win_name='all', v_image_count=(topk+1), display_mode=args.display_mode)

            cur_is_del = False
            topk_is_del = [False]*image_count
            is_exit = False
            face_num_str = ''
            while True:
                image = draw_image.copy()
                cur_bbox_idx = get_xy_bbox_idx(current_mouse_xy)
                if cur_bbox_idx is None:
                    cv2.rectangle(image, (0, img_height), (img_width*image_count-1, img_height*2-1), (0, 255, 0), 1)
                else:
                    bbox = image_bboxes[cur_bbox_idx]
                    cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)
                # print(cur_bbox_idx, current_mouse_xy)

                for select_idx in select_image_bbox_idxes:
                    bbox = image_bboxes[select_idx]
                    cv2.circle(image, (bbox[0]+14, bbox[1]+14), 8, (0, 0, 255), -1)

                if cur_is_del:
                    cv2.circle(image, (14, img_height * 2 - 14), 8, (0, 255, 0), -1)

                for _bbox_idx, is_del in enumerate(topk_is_del):
                    if is_del:
                        bbox = image_bboxes[_bbox_idx]
                        cv2.circle(image, (bbox[0]+14, bbox[3]-14), 8, (0, 255, 0), -1)

                for idx in range(0, image_count):
                    bbox = image_bboxes[idx]
                    score_str = '{:.03f}'.format(topk_score[idx])
                    cv2.putText(image, score_str, (bbox[0] + 2, bbox[3] - 4), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=2)
                cv2.putText(image, '%d/%d %d/%d' % (serial_num+1, face_count, xiaoqu_num, xiaoqu_count),
                            (8, 24), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), thickness=2)
                cv2.imshow('image', image)
                key = cv2.waitKey(50)
                if key == 27:
                    merge_change = merge_faceids_func(faceid, topk_faceids, select_image_bbox_idxes, faceid_file_feature_faceids, faceid_pairs)
                    select_image_bbox_idxes.clear()
                    select_change = merge_change
                    is_exit = True
                    break
                elif key == ord('d') or key == ord('D'):  # 下一个ID
                    merge_change = merge_faceids_func(faceid, topk_faceids, select_image_bbox_idxes, faceid_file_feature_faceids, faceid_pairs)
                    select_image_bbox_idxes.clear()
                    select_change = merge_change
                    break
                elif key == ord('c') or key == ord('C'):
                    cv2.destroyAllWindows()
                    cv2.namedWindow('image')
                    cv2.setMouseCallback('image', mouse_event)
                    continue
                elif ord('0') <= key <= ord('9'):  # 输入
                    if face_num_str == '':
                        face_num_str = chr(key)
                    else:
                        face_num_str += chr(key)
                    continue
                elif key == ord('f') or key == ord('F'):  # 播放视频
                    if cur_bbox_idx is None:
                        win_name = 'Pc'
                        image_file_path = topk_cur_file
                    else:
                        if cur_bbox_idx >= len(topk_files):
                            continue
                        win_name = 'P{}'.format(cur_bbox_idx)
                        image_file_path = topk_files[cur_bbox_idx]
                    print('F:', image_file_path)
                    video_name = os.path.dirname(image_file_path) + '.avi'
                    image_name = os.path.basename(image_file_path)
                    frame_num = int(image_name.split('_')[0])
                    play_video(video_name, frame_num, win_name=win_name)
                    continue
                elif key == ord('q') or key == ord('Q'):  # 显示所有人脸的所有图片
                    cv2.destroyAllWindows()
                    cv2.namedWindow('image')
                    cv2.setMouseCallback('image', mouse_event)
                    for i, face_id in enumerate(topk_faceids):
                        face_num_files = faceid_file_feature_faceids[face_id][0]
                        random.shuffle(face_num_files)
                        show_images(image_dir, face_num_files, win_name='{}'.format(i), display_mode=args.display_mode)
                    face_num_files = faceid_file_feature_faceids[faceid][0]
                    random.shuffle(face_num_files)
                    show_images(image_dir, face_num_files, win_name='cur', display_mode=args.display_mode)
                    continue
                elif key == ord('e') or key == ord('E'):  # 删除选中人脸
                    if cur_bbox_idx is None:
                        _face_id = faceid
                    else:
                        _face_id = topk_faceids[cur_bbox_idx]
                    _file_feature_faceids = faceid_file_feature_faceids[_face_id]
                    _file_feature_faceids[3] = -_file_feature_faceids[3]
                    if cur_bbox_idx is None:
                        cur_is_del = _file_feature_faceids[3] < 0
                    else:
                        topk_is_del[cur_bbox_idx] = _file_feature_faceids[3] < 0
                    continue
                elif key == ord('w') or key == ord('W'):  # 显示选中的人脸的所有图片ALL
                    if cur_bbox_idx is None:
                        _face_id = faceid
                        win_name = 'cur'
                    else:
                        _face_id = topk_faceids[cur_bbox_idx]
                        win_name = '{}'.format(cur_bbox_idx)
                    face_num_files = faceid_files.get(_face_id, None)
                    if face_num_files is not None:
                        random.shuffle(face_num_files)
                        show_images(image_dir, face_num_files, win_name=win_name, display_mode=args.display_mode)
                    continue
                elif key == ord('v') or key == ord('V'):  # 显示选中的人脸的所有图片
                    if cur_bbox_idx is None:
                        _face_id = faceid
                        win_name = 'cur'
                    else:
                        _face_id = topk_faceids[cur_bbox_idx]
                        win_name = '{}'.format(cur_bbox_idx)
                    face_num_files = faceid_file_feature_faceids[_face_id][0]
                    show_images(image_dir, face_num_files, win_name=win_name, display_mode=args.display_mode)
                    continue
                elif key == ord('s') or key == ord('S'):  # Save
                    if merge_change:
                        cur_xiaoqu_serial_num[xiaoqu] = serial_num + 1
                        merge_change = False
                    save_process(cur_xiaoqu_serial_num, merge_xiaoqu_faceid_file_feature_faceids, faceid_file_list_path, cur_id_idx_file_path)
                    start_save_time = time.time()
                    select_change = False
                    continue
                if args.use_temp:
                    # 定时保存
                    if time.time() - start_save_time >= save_time:
                        if select_change:
                            cur_xiaoqu_serial_num[xiaoqu] = serial_num + 1
                            select_change = False
                        save_process(cur_xiaoqu_serial_num, merge_xiaoqu_faceid_file_feature_faceids, tmp_select_file_list_path, tmp_cur_id_idx_file_path)
                        start_save_time = time.time()

            if is_exit:
                if merge_change:
                    cur_xiaoqu_serial_num[xiaoqu] = serial_num + 1
                save_process(cur_xiaoqu_serial_num, merge_xiaoqu_faceid_file_feature_faceids, faceid_file_list_path, cur_id_idx_file_path)
                return
        for faceid, file_feature_faceids in faceid_file_feature_faceids.items():
            if file_feature_faceids is None:
                continue
            file_feature_faceids[1] = None
        cur_xiaoqu_serial_num[xiaoqu] = face_count
    save_process(cur_xiaoqu_serial_num, merge_xiaoqu_faceid_file_feature_faceids, faceid_file_list_path, cur_id_idx_file_path)
    print('finish!')


if __name__ == '__main__':
    main()
